//==============================================================================
// SuffixTree.cc
// ----------------------------------------------------------------------------
//
// $Id: SuffixTree.cc,v 1.5 2000/03/22 14:14:33 kdurbin Exp kdurbin $
//------------------------------------------------------------------------------
// $Log: SuffixTree.cc,v $
// Revision 1.5  2000/03/22 14:14:33  kdurbin
// Discovered, and fixed, terrible slowdown due to using
// push_back on a ResultVec.   Various other speed optimizations
// implemented (inlining, etc.)
//
// Revision 1.4  2000/03/08 00:51:22  kdurbin
// Implemented own stack for DFS searching. Using recursion risked
// stack overflow.   Trailing special character was not being added
// to the source sequence.  Fixed that bug also.
//
// Revision 1.3  2000/03/02 20:27:47  kdurbin
// Restore from Mac after ISP connection failure.
// Don't think there are any differences....
//
// Revision 1.2  2000/02/01 00:57:46  kdurbin
// Minor touchups.
//
//
//==============================================================================
//.......1.........2.........3.........4.........5.........6.........7.........8
//

#include "SuffixTree.hh"

SuffixTreeEdge gBadEdge;

//---------------------------------------------------------------------------
// * IsBase
//---------------------------------------------------------------------------
bool
SuffixTree::IsBase(char base)
{
  switch (base)
	 {
	 case 'a': 
	 case 'c': 
	 case 'g': 
	 case 't': 
		return(true);
		break;
	 default:
		return(false);
	 }
}


//-------------------------------------------
// * SplitEdge
//-------------------------------------------
// Splits an edge like:
// 
// (originNode)--<EdgeIn>--(endNode)
//
// (originNode)--<EdgeIn>--(newNode)--<NewEdge>--(endNode)
//
// In a separate step, a new edge will be added to newNode 
// to represent the string that initiated the split.  
//
inline long 
SuffixTree::SplitEdge(SuffixTreeEdge *pEdgeIn,SuffixTreeSuffix &s)
{	
  // pEdgeIn is going to be transformed, so we do these computations first
  long LeafFirst = pEdgeIn->mFirstCharIdx+s.mLastCharIdx-s.mFirstCharIdx+1;
  long LeafLast = pEdgeIn->mLastCharIdx;
  long LeafEnd = pEdgeIn->mEndNode;
		
  // EdgeIn is transformed... removing characters for the suffix and 
  // giving it a new end node.  Note that there is no choice about 
  // this because pEdgeIn is tightly bound to s.mOriginNode!
  long NewLast = pEdgeIn->mFirstCharIdx+s.mLastCharIdx-s.mFirstCharIdx;
    //pEdgeIn->mFirstCharIdx=NewFirst; // First edge is unchanged.
  pEdgeIn->mLastCharIdx=NewLast;
  pEdgeIn->mEndNode = mNodeCount++;  // Introduces a new node. 
  mNodes[pEdgeIn->mEndNode].mSuffixNode = s.mOriginNode;

  //	Edge *pNewEdge = GetEdge(s.mOriginNode,mSource[pEdgeIn->mFirstCharIdx]);
  //pNewEdge->mFirstCharIdx=NewFirst;
  //pNewEdge->mLastCharIdx=NewLast;
  //pNewEdge->mStartNode = s.mOriginNode;
  //pNewEdge->mEndNode = mNodeCount++;  // Introduces a new node. 
  //mNodes[pNewEdge->mEndNode].mSuffixNode = s.mOriginNode;

  // This is a new edge only in the sense of having never been used before. 
  // Edges are pre-allocated, so GetEdge just returns a pointer to the 
  // appropriate pre-allocated edge.  No new memory is allocated here.
  SuffixTreeEdge *pNewLeafEdge = GetEdge(pEdgeIn->mEndNode,mSource[LeafFirst]);
  pNewLeafEdge->mFirstCharIdx = LeafFirst;
  pNewLeafEdge->mLastCharIdx = LeafLast;
  pNewLeafEdge->mEndNode = LeafEnd;
  
  return(pEdgeIn->mEndNode);
}

//-------------------------------------------
// * Canonize
//-------------------------------------------
// A suffix in the tree is denoted by a Suffix structure
// that denotes its last character.  The canonical
// representation of a suffix for this algorithm requires
// that the mOriginNode be the closest node to the end
// of the tree.  To force this to be true, we have to
// slide down every edge in our current path until we
// reach the final node.
void 
SuffixTree::CanonizeSuffix(SuffixTreeSuffix *pSuffix)
{
	if (!pSuffix->Explicit()){
	  // Get the edge corresponding to the suffix  
	  register SuffixTreeEdge *pEdge=GetEdge(pSuffix->mOriginNode,mSource[pSuffix->mFirstCharIdx]);
	  register long EdgeSpan = pEdge->mLastCharIdx - pEdge->mFirstCharIdx;

	  // If the edge does not cover all of the suffix, we will need to slide down...
	  while ( EdgeSpan <= (pSuffix->mLastCharIdx - pSuffix->mFirstCharIdx ) ) {

		 // Update suffix first char to point to the first character after the 
		 // previous edge.   
		 pSuffix->mFirstCharIdx = pSuffix->mFirstCharIdx + EdgeSpan + 1;
		 pSuffix->mOriginNode = pEdge->mEndNode;

		 if (pSuffix->mFirstCharIdx <= pSuffix->mLastCharIdx ) {
			pEdge = GetEdge(pEdge->mEndNode,mSource[pSuffix->mFirstCharIdx]);
			EdgeSpan = pEdge->mLastCharIdx - pEdge->mFirstCharIdx;
		 }

	  }
	}
}


//-------------------------------------------
// * EdgeIDValid
//-------------------------------------------
//
bool 
SuffixTree::EdgeIDValid(int ID)
{
	if ((ID>=0) && (ID < CHARACTER_SET_SIZE)) return(true);
	else return(false);

}

//----------------------------------------------------------------
// * NodeIDValid
//----------------------------------------------------------------
//
inline bool
SuffixTree::NodeIDValid(long NodeID)
{
 	if ((NodeID < mNodeCount) && (NodeID >=0)) return(true);
 	else return(false);
}

//----------------------------------------------------------------
// * EdgeValid
//----------------------------------------------------------------
//
inline bool
SuffixTree::EdgeValid(SuffixTreeEdge *pEdge)
{
	if (pEdge->mFirstCharIdx >=0) return(true);
	else return(false);
}


//-------------------------------------------
// * CharToIndex
//-------------------------------------------
//
inline int 
SuffixTree::CharToIndex(char c)
{
	switch(c){
	case 'a': return(0);break;
	case 'c': return(1);break;
	case 'g': return(2);break;
	case 't': return(3);break;
	case 'n': return(4);break;
	case '$': return(5);break;
	default: 
	  cout<<"Invalid char:("<<c;
	  int cVal = (int) c;
	  cout<<"),val="<<cVal<<" ."<<endl;
	  return(CHARACTER_SET_SIZE);break; // always means something went wrong. 
	}
}

//-------------------------------------------
// * GetEdge
//-------------------------------------------
// Gets an edge using character as an index. 
//
inline SuffixTreeEdge * 
SuffixTree::GetEdge(long nodeID,char edgeChar)
{		
	int EdgeID = CharToIndex(edgeChar);
	if (EdgeID != CHARACTER_SET_SIZE){
	  return(GetEdgeWithID(nodeID,EdgeID));
	}else{
	  gBadEdge.mFirstCharIdx = -1000;
	  return(&gBadEdge);
	}
}

//-------------------------------------------
// * GetEdgeWithID
//-------------------------------------------
// Returns a pointer to the edge at the specified node
// that starts with the given character.  It is important
// to realize that the edges are all pre-allocated, so 
// this function does not in any way allocate an edge. 
// Also, the edges are pre-associated with specific 
// nodes in memory.  Many subtle errors can arise from 
// acting as though this function creates edges.   
//
inline SuffixTreeEdge * 
SuffixTree::GetEdgeWithID(long nodeID,int EdgeID)
{
	if ((nodeID<0) || (nodeID >= mNodeCount)){
		return(&mInvalidEdge);
	}
		
	if (!EdgeIDValid(EdgeID)){
		return(&mInvalidEdge);
	}
	
	SuffixTreeEdge *pFoundEdge = &(mNodes[nodeID].mEdges[EdgeID]);
	return(pFoundEdge);
}

//-------------------------------------------
// * AddNewLeafEdge  
//-------------------------------------------
// KJD REALLY LEAF ALWAYS??????
// Create a new leaf edge... this includes creating the leaf node as
// well as inserting the new edge into the parent (newStartNode).  
// 
void 
SuffixTree::AddNewLeafEdge(long newFirstChar,long newLastChar,long newStartNode,
							long endSuffixLink)
{
	 mNodeCount++; // add node to the global list. 
	 long newEndNode = mNodeCount-1;
	 
	 // KJD this isn't needed if the edge really is always a leaf..... 
	 //mNodes[newEndNode].mSuffixNode = endSuffixLink;

	 // Each leaf is numbered sequentially, giving the position of this 
	 // suffix in the source. 
	 mNodes[newEndNode].mLeafIdx = mLeafNum++;  
		
	 SuffixTreeEdge *pNewEdge = GetEdge(newStartNode,mSource[newFirstChar]);
			
	 pNewEdge->mFirstCharIdx = newFirstChar;
	 pNewEdge->mLastCharIdx = newLastChar;
	 //pNewEdge->mStartNode = newStartNode;
	 pNewEdge->mEndNode = newEndNode;
}


//-------------------------------------------
// * PrintEdges
//-------------------------------------------
// Prints all the valid edges.  Very useful for 
// debugging.  
void 
SuffixTree::PrintEdges()
{
	for(long m = 0;m<=mMaxSourceIndex;m++)
		cout<<mSource[m];
	cout<<endl;

	cout<<"mNodeCount="<<mNodeCount<<endl;

    cout << "Start\tEnd\tLeaf#\tSuf\tFirst\tLast\tString\n";
    
    // For each node
    for(long i = 0;i<mNodeCount;i++){
    	int j = 0;
    	do{
    		SuffixTreeEdge *pS = &(mNodes[i].mEdges[j]);
    		//cout<<"Node["<<i<<"].Edge["<<j<<"].mStartNode="<<s.mStartNode<<endl;
    		//cout<<"first_char_idx="<<s.mFirstCharIdx<<endl;
    		if (EdgeValid(pS)){
    			cout << i <<"\t";
             cout	<< pS->mEndNode << "\t"
             	<< mNodes[ pS->mEndNode ].mLeafIdx<< "\t"
             	<< mNodes[ pS->mEndNode ].mSuffixNode << "\t"
             	<< pS->mFirstCharIdx << "\t"
             	<< pS->mLastCharIdx << "\t";
 			
        		for ( long k = pS->mFirstCharIdx ;k <= pS->mLastCharIdx;k++ )
            		cout << mSource[ k ];
        		cout << "\n";
    		}
    		j++;
    	}while(j<CHARACTER_SET_SIZE);
    }
}

//-------------------------------------------
// * MatchingEdgeExists
//-------------------------------------------
// Checks to see if there is already an edge for the given 
// suffix.   For explicit nodes, this means that there is an edge 
// coming out of the node with that character, for implicit nodes, 
// this means does the appropriate character in the edge match this 
// one...
// 
// 
inline bool
SuffixTree::MatchingEdgeExists(SuffixTreeSuffix &active,int LastCharIdx,SuffixTreeEdge* &pEdge)
{
  if (active.Explicit()){
	 pEdge = GetEdge(active.mOriginNode,mSource[LastCharIdx]);
	 if (EdgeValid(pEdge)) return(true); 
  }else{
	 pEdge = GetEdge(active.mOriginNode,mSource[active.mFirstCharIdx]);
	 long span = active.mLastCharIdx - active.mFirstCharIdx;
	 if ( mSource[ pEdge->mFirstCharIdx + span + 1 ] == mSource[ LastCharIdx ] ){
		return(true);
	 }
  }
  return(false);
}


//-------------------------------------------
// * AddPrefix
//-------------------------------------------
// This routine constitutes the heart of the algorithm.  It is called repeatedly, once 
// for each of the prefixes of the input string.  The prefix in question is denoted
// by the index of its last character.
//
// At each prefix, we start at the active point, and add a new edge denoting the new 
// last character, until we reach a point where the new edge is not needed due to
// the presence of an existing edge starting with the new last character.  This point is 
// the end point.
//
// Luckily for us, the end point just happens to be the active point for the next pass 
// through the tree.  All we have to do is update it's mLastCharIdx to indicate
// that it has grown by a single character, and then this routine can do all its work 
// one more time.
//
//
// Details :  Examine each suffix of the current prefix (identified by it's last 
// character... PrefixEndIdx) in order from largest to smallest (e.g.  
// mSource[0,PrefixEndIdx], mSource[1,PrefixEndIdx], mSource[2,PrefixEndIdx], etc.) 
// to see if that suffix is represented in the tree (i.e. there is a matching edge).   
// As soon as a matching edge for any suffix is found quit, because all smaller suffixes 
// will already be in the tree (by construction).   Until a matching suffix is found, however, 
// add edges for each suffix in succession.  Each such new edge will either be a brand new edge 
// on an explicit node, or a split edge on an implicit node.  After adding an edge, examine the 
// next smaller suffix to see if it is already in the tree.  Suffix-links are used to move 
// directly to the part of the tree where the next smaller suffix would be, if  it is there.   
// This is critical because starting from the root to search for a suffix in the tree makes 
// the construction algorithm non-linear.  
//
void 
SuffixTree::AddPrefix( SuffixTreeSuffix &active, long PrefixEndIdx )
{
  register long ParentNode;
  register long LastParentNode = -1;
	
  register SuffixTreeEdge *pEdge;
  ParentNode = active.mOriginNode;

  // If the edge is implicit, pEdge should exist, and will be returned 
  // by MatchingEdgeExists, it just may not have LastChar at the next 
  // location along the edge (as specified by active.mLastCharIdx+1).  If 
  // the node is explicit, there may or may not be an edge at all...
  while(!MatchingEdgeExists(active,PrefixEndIdx,pEdge)){

	 // No matching edge exists.  If implicit, this means that the edge is there 
	 // but the next character doesn't match, so we need to split the edge and 
	 // then add a leaf edge.   For explicit nodes, this means there simply isn't 
	 // an edge out of that node, so we just add a new leaf edge.   
	 if (active.Implicit()){
		ParentNode = SplitEdge(pEdge,active);
	 }
	 AddNewLeafEdge(PrefixEndIdx,mMaxSourceIndex,ParentNode,-1);
    
	 // Update suffix link, except when last parent was root...
	 if ( LastParentNode > 0 )
		mNodes[ LastParentNode ].mSuffixNode = ParentNode;      

	 LastParentNode = ParentNode;

	 // Move to next smaller suffix, following a suffix link unless we are at the 
	 // origin....
	 if ( active.mOriginNode == 0 ){
		active.mFirstCharIdx++;
	 }else{
		// Sanity test... should never happen. 
		if (mNodes[active.mOriginNode].mSuffixNode <0){
		  cout<<"ERROR: Attempt to follow a non-valid suffix link. \n";
		}
		active.mOriginNode = mNodes[ active.mOriginNode ].mSuffixNode;
	 }
	 CanonizeSuffix(&active);
	 ParentNode = active.mOriginNode;
  }

  // Update suffix-link if needed...
  if ( LastParentNode > 0 )
	 mNodes[ LastParentNode ].mSuffixNode = ParentNode;

  // Now endpoint is the next active point...
  active.mLastCharIdx++;  
  CanonizeSuffix(&active);
}



#if 0
//-------------------------------------------
// * OldAddPrefix
//-------------------------------------------
// This routine constitutes the heart of the algorithm.
// It is called repetitively, once for each of the prefixes
// of the input string.  The prefix in question is denoted
// by the index of its last character.
//
// At each prefix, we start at the active point, and add
// a new edge denoting the new last character, until we
// reach a point where the new edge is not needed due to
// the presence of an existing edge starting with the new
// last character.  This point is the end point.
//
// Luckily for us, the end point just happens to be the
// active point for the next pass through the tree.  All
// we have to do is update it's mLastCharIdx to indicate
// that it has grown by a single character, and then this
// routine can do all its work one more time.
//
void 
SuffixTree::OldAddPrefix( SuffixTreeSuffix &active, long LastCharIdx )
{
    register long ParentNode;
    register long LastParentNode = -1;
	
    for ( ; ; ) {
		register SuffixTreeEdge *pEdge;
      ParentNode = active.mOriginNode;

		// Step 1 is to try and find a matching edge for the given node.
		// If a matching edge exists, we are done adding edges, so we break
		// out of this big loop.
		if ( active.Explicit() ) {
			if (active.mOriginNode < 0){
        		cout<<"Eureka"<<endl;
        	}
        
   		pEdge = GetEdge(active.mOriginNode,mSource[LastCharIdx]);  	        
        	if (EdgeValid(pEdge)) break;

    	} else { //implicit node, a little more complicated
 			pEdge = GetEdge(active.mOriginNode,mSource[active.mFirstCharIdx]);
			long span = active.mLastCharIdx - active.mFirstCharIdx;
			cout<<"**Span="<<span<<endl;
			// Is there a matching "edge"?  If so, done with this iteration...
 			if ( mSource[ pEdge->mFirstCharIdx + span + 1 ] == mSource[ LastCharIdx ] )
				  break;

			// Otherwise, it's time for a split...
			ParentNode = SplitEdge(pEdge,active);
		}

		// If we didn't break out of this loop, we didn't find a matching edge, 
		// so we create a new one, add it to the tree at the parent node position, 
		// and insert it into the hash table.  When we create a new node, it also
		// means we need to create a suffix link to the new node from
		// the last node we visited.

    	// Create a new leaf edge... this includes creating the leaf node as
		// well as inserting the new edge into the parent (newStartNode).      
        AddNewLeafEdge(LastCharIdx,mMaxSourceIndex,ParentNode,-1);
          
        if ( LastParentNode > 0 )
            mNodes[ LastParentNode ].mSuffixNode = ParentNode;
      
        LastParentNode = ParentNode;

		// This final step is where we move to the next smaller suffix
        if ( active.mOriginNode == 0 ){
            active.mFirstCharIdx++;
        }else{
        	if (mNodes[active.mOriginNode].mSuffixNode <0){
        		cout<<"Attempt to follow a non-valid suffix link. \n";
        	}
        
            active.mOriginNode = mNodes[ active.mOriginNode ].mSuffixNode;
        }
        CanonizeSuffix(&active);
    }

    if ( LastParentNode > 0 )
        mNodes[ LastParentNode ].mSuffixNode = ParentNode;
    active.mLastCharIdx++;  //Now the endpoint is the next active point
    
    CanonizeSuffix(&active);

}
#endif

//-------------------------------------------
// * Create
//-------------------------------------------
// Creates a suffix tree for the given input. 
//
void
SuffixTree::Create(const char *source)
{
	 mSource = source;
	 mMaxSourceIndex = strlen(mSource) - 1;		
  	
	 // This allocation may look innocous enough... but in most cases it's 
	 // asking for a LOT of memory...
	
	 // Max leaves = mMaxSourceIndex;
	 // Worst case is a full tree, depth=log(CHARACTER_SET_SIZE,mMaxSourceIndex)
	 // containing CHARACTER_SET_SIZE^depth nodes = mMaxSourceIndex; 
	
	 long MaxNodes = 2*(mMaxSourceIndex+2);
	 cout<<"MaxNodes = "<<MaxNodes<<endl;
	 mNodes = new SuffixTreeNode[MaxNodes];
	 //cout<<MaxNodes<<" nodes allocated (";
	 cout<<MaxNodes*sizeof(SuffixTreeNode)<<") bytes ..."<<flush;
  	
	 // The active point is the first non-leaf suffix in the
	 // tree.  We start by setting this to be the empty string
	 // at node 0.  The AddPrefix() function will update this
	 // value after every new prefix is added.
	 SuffixTreeSuffix active( 0, 0, -1 );  // The initial active prefix

	 for ( long prefixNum = 0 ; prefixNum <= mMaxSourceIndex ; prefixNum++ ){
		//if (!(i%50000))
		//cout<<"------------"<<endl;
		//cout<<"Adding prefix "<<prefixNum<<endl;
		//cout<<"Active: ";
		//active.Print();
		//cout<<endl;
		AddPrefix( active, prefixNum );
		//PrintEdges();
	 }

	 // Pre allocate worst case space, plus a little ;-) for the stack 
	 // that we use for searching. 
	 mpStack = new STreeSearchStack(mNodeCount+2); 
	 mpStack->clearVisited(); // Redundant, but making sure.
}

//----------------------------------------------------------------
// * StackChildren
//----------------------------------------------------------------
// Adds the child nodes of NodeID to the stack DFSStack
// 
inline void
SuffixTree::StackChildren(long NodeID,STreeSearchStack *pStack,SuffixTreeEdge *pEdge)
{
  register int edgeID;

  // Each edge radiating out of node NodeID will end in some 
  // node, stored in the edge structure as mEndNode.
  for(edgeID = 0;edgeID < CHARACTER_SET_SIZE;edgeID++){

	 //pEdge = GetEdgeWithID(NodeID,edgeID);
	 pEdge = &(mNodes[NodeID].mEdges[edgeID]);

	 // Edges are pre-allocated.  Unused edges are simply marked invalid
	 // and must be skipped. If CHARACTER_SET_SIZE is large, there are 
	 // more efficient schemes.

	 // if (EdgeValid(pEdge)){
	 if(pEdge->mFirstCharIdx >=0){
		pStack->push(pEdge->mEndNode);
	 }
  }
}

//----------------------------------------------------------------
// * DFSLeafIdx
//----------------------------------------------------------------
// Performs a depth first search from the given node, returning 
// the indexes of the leaves below that node which, by construction,
// give the locations in the source where the prefix denoted by the 
// node occur.  This is the core of finding a query in the tree so 
// it is important that it be as fast as possible.  Additional easy 
// speed-up can be had by pre-allocating Indexes rather than using push_back.    
// 
// Algorithmic note:  Although this does a DFS, the running time of 
// search for a query is independent of the source (tree) size.  
// This DFS is only down the sub-tree denoted by the given prefix, and
// every leaf in this subtree marks an occurence of this prefix in the
// source sequence.   Since every node has at least two edges, the number
// of edges in this subtree is proportional to the number of leaves, thus
// the running time of query is proportional to the size of the 
// query (in FindFirst()) and the number of occurences of the query in 
// the source.    
// 
void
SuffixTree::DFSLeafIdx(long NodeID,vector<long> &Indexes)
{ 
  register SuffixTreeEdge *pEdge;

  //mpStack->clearVisited(); KJD Why does it fail if I don't do this???

  // If this node is a leaf, we're done, record the index and split...
  if (mNodes[NodeID].mLeafIdx >=0){
	 Indexes.resize(Indexes.size()+1);
	 Indexes[Indexes.size()-1] = mNodes[NodeID].mLeafIdx;
	 //Indexes.push_back(mNodes[NodeID].mLeafIdx);
	 return;
  }else{
	 // Otherwise add the children of this node to the stack...
	 StackChildren(NodeID,mpStack,pEdge);
  }

  // Until the stack is empty, search each node.  
  while(mpStack->size()>0){
	 register long CNode = mpStack->pop();
	 
	 // Debug test for extra work. Since this is a tree, DFS should never
	 // see the same node twice.  If we do, the tree has been badly constructed.
	 //if (mpStack->visited(CNode)) cout<<"########## Visited "<<CNode<<endl;
	 //mpStack->visit(CNode);
			
	 // If this node is a leaf, record the index and go to the next node 
	 // in the stack (i.e. terminate the DFS down this branch).  
	 if (mNodes[CNode].mLeafIdx >=0){	 
		Indexes.resize(Indexes.size()+1);
		Indexes[Indexes.size()-1] = mNodes[CNode].mLeafIdx;
		//Indexes.push_back(mNodes[CNode].mLeafIdx);
	 }else{
		StackChildren(CNode,mpStack,pEdge);
	 }
  };
}

//----------------------------------------------------------------
// * Find
//----------------------------------------------------------------
//
void 
SuffixTree::FindAll(const char * query, long length,vector<long> &Locations)
{
	 long FirstLocation;
	 long FirstNodeID;
	 FirstNodeID = FindFirst(query,length,FirstLocation);
	 
	 // No occurences found
	 if (FirstLocation < 0){
		return;
	 }
	 DFSLeafIdx(FirstNodeID,Locations);
	
	 // When there is only one occurence of a substring, there will be 
	 // no leaf nodes corresponding to the string, i.e. it ends in a point
	 // internal to an edge.  In this case, simply return the single 
	 // location found. 
	 if (Locations.size() == 0) Locations.push_back(FirstLocation);
}

#if 0
//----------------------------------------------------------------
// * FindFirst
//----------------------------------------------------------------
// Finds the first occurence of query in the suffix tree.  Returns
// the location and the last node of this match, which is used 
// by find all to recover the remainder of the hits.  
//  
// 
long
SuffixTree::FindFirst(const char *query,long length,long &Location)
{	
  register long queryOffset = 0;
  register long Length = length;
  register long sourceIdx=0;
	  
  register SuffixTreeEdge *pEdge = GetEdge(0,query[0]); // KJD Origin should be zero. 
  if (!EdgeValid(pEdge)){
	 Location = -1;
	 return(pEdge->mEndNode);
  }


  do { 
	 sourceIdx = pEdge->mFirstCharIdx;
	 register long LastCharIdx = pEdge->mLastCharIdx;

	 while((mSource[sourceIdx] == query[queryOffset]) &&
	 		  (queryOffset < Length) &&
	 		 (sourceIdx <= LastCharIdx)){
	 	queryOffset++;
	 	sourceIdx++;
	 };

	 // Three things could happen to exit the loop:
	 // 1. we've run out of query(found match)->Location=source-L+1, node = edge->end
	 // 2. mismatch -> in which case we are done, we can exit this whole routine. 
	 // 3. we've reached the end of this node-> get next edge and continue
	 if (queryOffset >= Length){
		Location = sourceIdx - Length+1;
		return(pEdge->mEndNode);
	 }else{
		if (sourceIdx > LastCharIdx){
		  pEdge = GetEdge(pEdge->mEndNode,query[queryOffset]);
		}else{
		  if (mSource[sourceIdx] != query[queryOffset]){
			 Location = -1;
			 return(pEdge->mEndNode);
		  }
		}
	 }
  }while(EdgeValid(pEdge));
  
  cout<<"End due to invalid edge.\n";
  Location = -1;
  return(pEdge->mEndNode);
}
#endif

//----------------------------------------------------------------
// * FindPrefix
//----------------------------------------------------------------
// Finds the first occurence of a prefix of query in the suffix tree.  
// Returns the location and the last node of this match
//  
// 
long
SuffixTree::FindPrefix(const char *query,long length,long &Location)
{	
	 register long queryOffset = 0,edgeOffset=0;
	 register SuffixTreeEdge *pEdge = GetEdge(0,query[0]); // KJD Origin should be zero. 
	 register SuffixTreeEdge *pLastEdge = pEdge;
	 register long sourceIdx=pEdge->mFirstCharIdx;
	 long lastSourceIdx = sourceIdx;
	 // These bool conditions are mostly to improve readability of the loop tests.
	 register bool bMismatch=false;
	 register bool bIsLeaf = false;
	 register bool bQueryEnd = false;
	 register bool bSourceContainedInEdge = true;
	 long NodeOut = 0;

	 // This can happen if no suffix of the source starts with the same 
	 // character as the first character of the query.   None of the prefix 
	 // is in the tree, so the starting node is 0 and the starting char is 0. 
	 if (!EdgeValid(pEdge)){
		Location = 0;
		return(0);
	 }

	 do { 
		do{			
		  // Get the next character and see if it matches the query or if one of 
		  // the query or source is an n.  KJD: Note, changed test to IsBase
		  // in order to treat ALL non bases as though they were n's (i.e. 
		  // those pesky K,Y, and other equivalence classes).  
		  if ((mSource[sourceIdx] != query[queryOffset])){
		  //	  && (mSource[sourceIdx]!='n') && (query[queryOffset]!='n')){
		  //				IsBase(mSource[sourceIdx]) && IsBase(query[queryOffset])){
			 bMismatch = true;
		  }else{
			 queryOffset++;
			 
			 // // We've exhausted the query, so we're done. 
			 if (queryOffset > (length-1)) {
				bQueryEnd = true;
			 }
			 else {
				edgeOffset++;
				lastSourceIdx = sourceIdx;
				sourceIdx = pEdge->mFirstCharIdx + edgeOffset;
			 } 
		  }
		  bSourceContainedInEdge = (sourceIdx <= pEdge->mLastCharIdx);
		}while(!bMismatch && bSourceContainedInEdge && !bQueryEnd);
		  
		// Loop ended because it reached end of edge, not because of a mismatch.
		// or query end.  
		if ((!bSourceContainedInEdge) && (!bQueryEnd) && (!bMismatch)) {
		  pLastEdge = pEdge;
		  pEdge = GetEdge(pEdge->mEndNode,query[queryOffset]);
		  edgeOffset=0;
		  sourceIdx = pEdge->mFirstCharIdx;
		}
	 }while(EdgeValid(pEdge) && !bMismatch && !bQueryEnd);
	
	 // Only matched a fragment of 
	 if (bMismatch) {
		Location = sourceIdx;
		NodeOut = pLastEdge->mEndNode;
	 }else { 
		//sourceIdx points to the end of the query.  We want to know where the 
		// beginning is. 
		Location = sourceIdx;
		NodeOut = pLastEdge->mEndNode;
	 }
	 // This happens when the prefix mismatches precisely at a node. 
	 if (Location == -1) Location = 0;
	 return(NodeOut);
}


//----------------------------------------------------------------
// * FindFirst
//----------------------------------------------------------------
// Finds the first occurence of query in the suffix tree.  Returns
// the location and the last node of this match, which is used 
// by find all to recover the remainder of the hits.  
//  
// 
long
SuffixTree::FindFirst(const char *query,long length,long &Location)
{	
	 register long queryOffset = 0,edgeOffset=0;
	 register SuffixTreeEdge *pEdge = GetEdge(0,query[0]); // KJD Origin should be zero. 
	 register long sourceIdx=pEdge->mFirstCharIdx;
	  
	 // These bool conditions are mostly to improve readability of the loop tests.
	 register bool bMismatch=false;
	 register bool bIsLeaf = false;
	 register bool bQueryEnd = false;
	 register bool bSourceContainedInEdge = true;
	 
	 // Sanity check 
	 if (!EdgeValid(pEdge)) return(-1);
	 
	 do { 
		do{			
		  // Get the next character and see if it matches the query or if one of 
		  // the query or source is an n.  KJD: Note, changed test to IsBase
		  // in order to treat ALL non bases as though they were n's (i.e. 
		  // those pesky K,Y, and other equivalence classes).  
		  if ((mSource[sourceIdx] != query[queryOffset])){
		  //	  && (mSource[sourceIdx]!='n') && (query[queryOffset]!='n')){
		  //				IsBase(mSource[sourceIdx]) && IsBase(query[queryOffset])){
			 bMismatch = true;
		  }else{
			 queryOffset++;
			 
			 // // We've exhausted the query, so we're done. 
			 if (queryOffset > (length-1)) {
				bQueryEnd = true;
			 }
			 else {
				edgeOffset++;
				sourceIdx = pEdge->mFirstCharIdx + edgeOffset;
			 } 
		  }
		  bSourceContainedInEdge = (sourceIdx <= pEdge->mLastCharIdx);
		}while(!bMismatch && bSourceContainedInEdge && !bQueryEnd);
		  
		// Loop ended because it reached end of edge, not because of a mismatch.
		// or query end.  
		if ((!bSourceContainedInEdge) && (!bQueryEnd) && (!bMismatch)) {
		  pEdge = GetEdge(pEdge->mEndNode,query[queryOffset]);
		  edgeOffset=0;
		  sourceIdx = pEdge->mFirstCharIdx;
		}
	 }while(EdgeValid(pEdge) && !bMismatch && !bQueryEnd);
	
	 
	 if (bMismatch) {
		Location = -1;
	 }
	 else 
	 { 
		//sourceIdx points to the end of the query.  We want to know where the 
		// beginning is. 
		Location = sourceIdx - length + 1;
	 }
	 
	 return(pEdge->mEndNode);
}


//----------------------------------------------------------------
// * Find
//----------------------------------------------------------------
//
long
SuffixTree::Find(const char *query,long length)
{	
	 long Location;
	 
	 FindFirst(query,length,Location);
	 return(Location);
	 
}
 
//-------------------------------------------
// * CreateST
//-------------------------------------------
//
void
SuffixTreeSearchEngine::CreateST(Sequence &Sequence)
{
  Initialize();

  // Create a suffix tree for this database sequence
  char * pSourceCString = new char[Sequence.size()+3];
  Sequence.ToLowerCase(); 
  Sequence.CopyToCString(pSourceCString);
  //cout<<Sequence<<endl;
  //cout<<pSourceCString<<endl;

  //long VerifyN = strlen(pSourceCString);
  //cout<<"Creating suffix tree ("<<Sequence.size();
  //cout<<" bp source)..."<<flush;
  
  // Every source must be terminated by a unique character. 
  pSourceCString[Sequence.size()]='$';
  pSourceCString[Sequence.size()+1]='\0';
  Create(pSourceCString);
  //cout<<" done."<<endl;
}

//-------------------------------------------
// * RemoveST
//-------------------------------------------
//
void
SuffixTreeSearchEngine::RemoveST()
{
  Clear();
}

//-------------------------------------------
// * CreateST
//-------------------------------------------
//
void
SuffixTreeSearchEngineConcatenated::CreateST(SequenceVector &SeqVec,
															int sstart,int send)
{
  Initialize();

  mFragmentBoundaries.clear();
  mSubjectsStartIdx = sstart;

  long BoundaryLoc = 0;

  // Save the locations of the boundaries...
  for(int i = sstart;i<=send;i++){
	 mConcatenatedSource.Concatenate(SeqVec[i]);
	 BoundaryLoc+=SeqVec[i].size();
	 mFragmentBoundaries.push_back(BoundaryLoc-1);
  }

  //cout<<"Concatenated source: "<<endl;
  //mConcatenatedSource.Print();
  //cout<<"Boundaries:"<<endl;
  //for(int i = 0;i< mFragmentBoundaries.size();i++){
  // cout<<mFragmentBoundaries[i]<<endl;
  //}

  // Create a suffix tree for this database sequence
  cout<<"mConcatenatedSource.size()="<<mConcatenatedSource.size()<<" "<<flush;
  char * pSourceCString = new char[mConcatenatedSource.size()+3];
  mConcatenatedSource.ToLowerCase(); 
  mConcatenatedSource.CopyToCString(pSourceCString);

  // Every source must be terminated by a unique character. 
  pSourceCString[mConcatenatedSource.size()]='$';
  pSourceCString[mConcatenatedSource.size()+1]='\0';
  cout<<"..create..."<<flush;
  Create(pSourceCString);
  
}

//-------------------------------------------
// * Search
//-------------------------------------------
//
// @todo Names are no longer stored in SearchResult
//       Several places previously extracted information out of SubjectName, 
//       such as sense and query location. Affected classes include SearchResult, SuffixTree
//
// @todo SuffixTree search currently saves non-hits as well as hits in SearchResult.
//       The old version of overlapper required this.  Does the new version?  It 
//       should be changed, in any case.    
//
// @todo Who uses mSubjects in SearchResult?
//
void
SuffixTreeSearchEngine::Search(SequenceVector &Queries,
										 SearchResultVec &Results)
{
  char * pQueryCString = new char[Queries[0].size()+2];

  // Search each query sequence against the suffix tree. 
  for(int qNum = 0;qNum < Queries.size();qNum++){

	 // Only re-allocate when forced to.  Since in most of my applications
	 // the queries are all the same size, this almost never happens. 
	 int QuerySize = Queries[qNum].size()+2;
	 if (QuerySize > Queries[0].size()){
		if (pQueryCString != NULL) delete pQueryCString;
		pQueryCString = new char[QuerySize+2];
	 }

	 Queries[qNum].CopyToCString(pQueryCString); 
	 
	 // Finds all the occurences of the query in the source. 
	 vector<long> Locations;
	 FindAll(pQueryCString,strlen(pQueryCString),Locations);

	 // KJD SPEED This section roughly doubles the search time. 
	 // Without this section, 130,000 25bp samples can be searched in 
	 // a second.  With this section, only 45-54,000 samples.   

	 // Amazingly, this section is the speed limiting factor
	 // for this search.   Results.push_back(tempResult) is very slow 
	 // compared to the code that's here now.  Obviously, much faster things 
	 // exist and could/should probably be used.   SearchResultVec either 
	 // needs to be rewritten to be more lightweight, or some stripped down 
	 // version needs to be put here. 

	 
	 // Record queries which have no hits. 
	 if (Locations.size() == 0){
		// KJD Why save these non-hits?  Because the subject/query lists
		// are built by scanning the results.  No results, no subject/query
		// in list.  It would be better to handle this at a high level
		// rather than here, but this works. 

		Results.resize(Results.size()+1);
		int idx = Results.size()-1;

		// Give unused fields some benign, or at least easily identfiable values
		//Results[idx].mScoreBits = 0;
		//Results[idx].mExpect = 0;
		Results[idx].mQueryLocation = 0;

		Results[idx].mSubjectLocation = -1;
		//Results[idx].mIdentitiesDen = Queries[qNum].size();
		//Results[idx].mIdentitiesNum = 0;
		Results[idx].mQueryIdx = qNum;
		//Results[idx].mQueryName = Queries[qNum].mDescription;

		// Often, original query location is encoded in the name.  If so, 
		// this value is extracted here into mQueryLocation.  
		//Results[idx].ExtractQueryLocFromName(Queries[qNum].mDescription,
		//												 Results[idx].mQueryLocation);

	 }else{
		register int idx = Results.size() -1;
		Results.resize(Results.size()+Locations.size());
		for(int lnum=0;lnum < Locations.size();lnum++){
		  idx++;
		  
		  // Give unused fields some benign, or at least easily identifiable values
		  //Results[idx].mScoreBits = 0;
		  //Results[idx].mExpect = 0;
		  Results[idx].mQueryLocation = 0;
		  Results[idx].mSubjectLocation = Locations[lnum];
		  //Results[idx].mIdentitiesDen = Queries[qNum].size();
		  //Results[idx].mIdentitiesNum = Queries[qNum].size();
		  Results[idx].mQueryIdx = qNum;
		  //Results[idx].ExtractQueryLocFromName(Queries[qNum].mDescription,
		  //													Results[idx].mQueryLocation);


		  //Results[startIdx+lnum].mQueryName = Queries[qNum].mDescription;
		  //Results[startIdx+lnum].ExtractQueryLocFromName();
		}
	 }

	 // KJD End Speed killing section -----------------------------------------------
  }
  //Results.GetSubjectList(Results.mSubjects); KJD Who uses subject list mSubjects?
  //Results.GetQueryList(Results.mQueries);
  delete pQueryCString;
}



//-------------------------------------------
// * CreateST
//-------------------------------------------
//
void
SuffixTreeSearchEngineConcatenated::CreateST(SequenceVector &SeqVec)
{
  CreateST(SeqVec,0,SeqVec.size()-1);
}

//-------------------------------------------
// * Synthetic
//-------------------------------------------
// Tests whether a hit is a synthetic hit (on the boundary of 
// two concatenated sequences.   This should be rare.  If 
// it's rare enough, this test can be eliminated. 
//
bool
SuffixTreeSearchEngineConcatenated::Synthetic(long Loc,int qsize)
{
  vector<long>::iterator nextGreater,nextLesser;
  nextGreater = lower_bound(mFragmentBoundaries.begin(),
									 mFragmentBoundaries.end(),Loc);
  nextLesser = nextGreater-1;

  long Diff = 0;
  vector<long>::iterator closest;
  // Loc is smaller than any location in the list
  if (nextGreater == mFragmentBoundaries.begin()){
	 closest = nextGreater;
	 //Diff = *nextGreater - Loc;
  }else{
	 // Loc is greater than any location in the list
	 if (nextGreater == mFragmentBoundaries.end()){
		closest = nextLesser;
		//Diff = Loc - *nextLesser;
	 }else{
		// Loc has greater and lesser value in the list
		// Return the smaller difference
		long GDiff = *nextGreater -Loc;
		long LDiff = Loc - *nextLesser;
		if (GDiff < LDiff) closest = nextGreater;
		else closest =nextLesser;
	 }
  }
  
  Diff = *closest - Loc;
  if ((Diff > 0) && (Diff < qsize)){
	 mSyntheticHitCount++;
	 return(true);
  }else{
	 return(false);
  }
}


//-------------------------------------------
// * ExtractSourceInfo
//-------------------------------------------
// Doing this is a big performance hit.  This will not be necessary when 
// the generalized stree code is working.   Also, this can be simplified
// if the source sizes are all the same. 
//
//
void
SuffixTreeSearchEngineConcatenated::ExtractSourceInfo(long cIdx,int &sIdx,long &sLoc)
{
  vector<long>::iterator nextGreater,nextLesser,temp;
  nextGreater = lower_bound(mFragmentBoundaries.begin(),
									 mFragmentBoundaries.end(),cIdx);

  sIdx = nextGreater - mFragmentBoundaries.begin();
  if (nextGreater == mFragmentBoundaries.begin()){
	 sLoc = cIdx;
  }else{
	 nextLesser = nextGreater-1;
	 sLoc = cIdx - *nextLesser - 1;
  }
}

//-------------------------------------------
// * Search
//-------------------------------------------
//
void
SuffixTreeSearchEngineConcatenated::Search(SequenceVector &Queries,
														 SearchResultVec &Results)
{

  mSyntheticHitCount = 0;
  char * pQueryCString = new char[Queries[0].size()+2];
  
  // Search each query sequence against the suffix tree. 
  for(int qNum = 0;qNum < Queries.size();qNum++){
	 if (Queries[qNum].size()>0){

		// Only re-allocate when forced to.  Since in most of my applications
		// the queries are all the same size, this almost never happens. 
		int QuerySize = Queries[qNum].size()+2;
		if (QuerySize > Queries[0].size()){
		  if (pQueryCString != NULL) delete pQueryCString;
		  pQueryCString = new char[QuerySize+2];
		}
	 
		Queries[qNum].CopyToCString(pQueryCString); 
		
		// Finds all the occurences of the query in the source. 
		vector<long> Locations;
		FindAll(pQueryCString,strlen(pQueryCString),Locations);

		//cout<<"Locations = "<<endl;
		//for(int i = 0;i<Locations.size();i++){
		//	cout<<Locations[i]<<endl;
		//}

		register int idx;
		// KJD SPEED : Amazingly, this section is the speed limiting factor
		// for this search.   Results.push_back(tempResult) is very slow 
		// compared to the code that's here now.  Obviously, much faster things 
		// exist and could/should probably be used.   SearchResultVec either 
		// needs to be rewritten to be more lightweight, or some stripped down 
		// version needs to be put here. 
		
		// Record queries which have no hits. 
		//if (Locations.size() == 0){
		// KJD Why save these non-hits?  Because the subject/query lists
		// are built by scanning the results.  No results, no subject/query
		// in list.  It would be better to handle this at a high level
		// rather than here, but this works. 
		
		//Results.resize(Results.size()+1);
		//idx = Results.size()-1;
		
		// Give unused fields some benign, or at least easily identfiable values
		//Results[idx].mScoreBits = 0;
		//Results[idx].mExpect = 0;
		//Results[idx].mQueryLocation = 0;
		
		//Results[idx].mSubjectLocation = -1;
		//Results[idx].mIdentitiesDen = Queries[qNum].size();
		//Results[idx].mIdentitiesNum = 0;
		//Results[idx].mQueryIdx = qNum;
		//Results[idx].mQueryName = Queries[qNum].mDescription;
		
		// Often, original query location is encoded in the name.  If so, 
		// this value is extracted here into mQueryLocation.  
		//Results[idx].ExtractQueryLocFromName(Queries[qNum].mDescription,
		//												 Results[idx].mQueryLocation);

		// }else{
		if (Locations.size() > 0){
		  if (mbFilterSyntheticHits){
			 idx = Results.size() -1;
			 for(int lnum=0;lnum < Locations.size();lnum++){
				if (!Synthetic(Locations[lnum],Queries[0].size())){
				  Results.resize(Results.size()+1);
				  idx++;
			 
				  // Give unused fields some benign, or at least easily 
				  // identifiable values
				  //Results[idx].mScoreBits = 0;
				  //Results[idx].mExpect = 0;
				  Results[idx].mQueryLocation = 0;
				  long sLoc;
				  int sIdx;
				  ExtractSourceInfo(Locations[lnum],sIdx,sLoc);
				  Results[idx].mSubjectLocation = sLoc;
				  // Translate into coordinates of input SequenceVec
				  Results[idx].mSubjectIdx = sIdx;
				
				  // Sanity test
				  if (sIdx >= mFragmentBoundaries.size()){
					 cout<<"FragmentBoundaries.size()="<<mFragmentBoundaries.size()<<endl;
					 cout<<"sLoc="<<sLoc<<endl;
					 cout<<"sIdx="<<sIdx<<endl;
					 cout<<"Queries size="<<Queries.size()<<endl;
					 exit(1);
				  }


				  Results[idx].mQueryIdx = qNum;
				  //Results[idx].mIdentitiesDen = Queries[qNum].size();
				  //Results[idx].mIdentitiesNum = Queries[qNum].size();
				  //Results[idx].ExtractQueryLocFromName(Queries[qNum].mDescription,
				  //											Results[idx].mQueryLocation);
				
				
				  //Results[startIdx+lnum].mQueryName = Queries[qNum].mDescription;
				  //Results[startIdx+lnum].ExtractQueryLocFromName();
				}
			 }
		  }else{
			 idx = Results.size() -1;
			 Results.resize(Results.size()+Locations.size());
			 for(int lnum=0;lnum < Locations.size();lnum++){
				idx++;
		  
				// Give unused fields some benign, or at least easily identifiable values
				//Results[idx].mScoreBits = 0;
				//Results[idx].mExpect = 0;
				Results[idx].mQueryLocation = 0;
				long sLoc;
				int sIdx;
				ExtractSourceInfo(Locations[lnum],sIdx,sLoc);
				Results[idx].mSubjectLocation = sLoc;
				Results[idx].mSubjectIdx = sIdx;
				
				// Sanity test
				if (sIdx >= mFragmentBoundaries.size()){
				  cout<<"FragmentBoundaries.size()="<<mFragmentBoundaries.size()<<endl;
				  cout<<"sLoc="<<sLoc<<endl;
				  cout<<"sIdx="<<sIdx<<endl;
				  cout<<"Queries size="<<Queries.size()<<endl;
				  exit(1);
				  // KJD BUG: This condition occurs sometimes when there is only 
				  // one sequence??  Basically if num sequences at a time includes
				  // all of the sequences, this doesn't happen, but sometimes when 
				  // number of sequences is > SequencesAtOnce, this condition occurs. 
				}
				
				Results[idx].mQueryIdx = qNum;
				//Results[idx].mIdentitiesDen = Queries[qNum].size();
				//Results[idx].mIdentitiesNum = Queries[qNum].size();
				//Results[idx].ExtractQueryLocFromName(Queries[qNum].mDescription,
				//												 Results[idx].mQueryLocation);
				
				//Results[startIdx+lnum].mQueryName = Queries[qNum].mDescription;
				//Results[startIdx+lnum].ExtractQueryLocFromName();
			 }
		  }
		  // KJD End Speed killing section -------------------------------------
		}
	 }
  }
  delete pQueryCString;
  //cout<<"Results.size()="<<Results.size()<<endl;
  //Results.GetSubjectList(Results.mSubjects);
  //Results.GetQueryList(Results.mQueries);
  //cout<<"SyntheticHitCount="<<mSyntheticHitCount<<endl;
}

